from depthaware.models.model_utils import *
from depthaware.models.ops.depthconv.module import DepthConv
from depthaware.models.ops.depthavgpooling.module import DepthAvgPooling
from depthaware.utils.utils import *


import torch.nn.functional as F
import torch
import torch.nn as nn
import torch.utils.model_zoo as model_zoo
from math import sqrt



class ConvModule(nn.Module):

    def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=1, dilation=1,
                 bn=False,
                 maxpool=False, pool_kernel=3, pool_stride=2, pool_pad=1, pool_ceil=False):
        super(ConvModule, self).__init__()
        conv2d = nn.Conv2d(inplanes, planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation)
        layers = []
        if bn:
            layers += [nn.BatchNorm2d(planes), nn.ReLU(inplace=True)]
        else:
            layers += [nn.ReLU(inplace=True)]
        if maxpool:
            layers += [nn.MaxPool2d(kernel_size=pool_kernel, stride=pool_stride, padding=pool_pad, ceil_mode=pool_ceil)]

        self.layers = nn.Sequential(*([conv2d] + layers))

    def forward(self, x):
        # x = self.conv2d(x)
        x = self.layers(x)
        return x


class DepthConvModule(nn.Module):

    def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=1, dilation=1, bn=False):
        super(DepthConvModule, self).__init__()

        conv2d = DepthConv(inplanes, planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation)
        layers = []
        if bn:
            layers += [nn.BatchNorm2d(planes), nn.ReLU(inplace=True)]
        else:
            layers += [nn.ReLU(inplace=True)]
        self.layers = nn.Sequential(*([conv2d] + layers))  # (*layers)

    def forward(self, x, depth):

        for im, module in enumerate(self.layers._modules.values()):
            if im == 0:
                x = module(x, depth)
            else:
                x = module(x)
        # x = self.conv2d(x, depth)
        # x = self.layers(x)
        return x


class VGG_layer(nn.Module):
    def __init__(self, batch_norm=False, depthconv=False):  # bn指是否归一化，True则nn.BatchNorm2d，否则nn.ReLU
        super(VGG_layer, self).__init__()
        in_channels = 3
        self.depthconv = depthconv
        if self.depthconv:
            self.conv1_1_depthconvweight = 1.  # nn.Parameter(torch.ones(1))
            # 卷积层输出尺寸： o = ⌊(i + 2p - k) / s⌋ + 1,
            self.conv1_1 = DepthConvModule(3, 64, bn=batch_norm)  # 卷积下采样kernel_size=3, stride=1, padding=1尺寸不变
        else:
            self.conv1_1 = ConvModule(3, 64, bn=batch_norm)
        # pool_kernel=3, pool_stride=2, pool_pad=1 与默认的VGG不同，池化层输出尺寸： o = ⌊(i + 2p - k) / s⌋ + 1,
        self.conv1_2 = ConvModule(64, 64, bn=batch_norm,
                                  maxpool=True, pool_kernel=2, pool_stride=2, pool_pad=0)

        if self.depthconv:
            self.conv2_1_depthconvweight = 1.  # nn.Parameter(torch.ones(1))
            # pool_kernel=3, pool_stride=2, pool_pad=1 与上面的池化同步
            self.downsample_depth2_1 = nn.AvgPool2d(2, padding=0, stride=2)
            self.conv2_1 = DepthConvModule(64, 128, bn=batch_norm)
        else:
            self.conv2_1 = ConvModule(64, 128, bn=batch_norm)
        self.conv2_2 = ConvModule(128, 128, bn=batch_norm,
                                  maxpool=True, pool_kernel=2, pool_stride=2, pool_pad=0)

        if self.depthconv:
            self.conv3_1_depthconvweight = 1.  # nn.Parameter(torch.ones(1))
            self.downsample_depth3_1 = nn.AvgPool2d(2, padding=0, stride=2)
            self.conv3_1 = DepthConvModule(128, 256, bn=batch_norm)
        else:
            self.conv3_1 = ConvModule(128, 256, bn=batch_norm)
        self.conv3_2 = ConvModule(256, 256, bn=batch_norm)
        self.conv3_3 = ConvModule(256, 256, bn=batch_norm,
                                  maxpool=True, pool_kernel=2, pool_stride=2, pool_pad=0, pool_ceil=True)

        if self.depthconv:
            self.conv4_1_depthconvweight = 1.  # nn.Parameter(torch.ones(1))
            self.downsample_depth4_1 = nn.AvgPool2d(2, padding=0, stride=2, ceil_mode=True)
            self.conv4_1 = DepthConvModule(256, 512, bn=batch_norm)
        else:
            self.conv4_1 = ConvModule(256, 512, bn=batch_norm)
        self.conv4_2 = ConvModule(512, 512, bn=batch_norm)
        self.conv4_3 = ConvModule(512, 512, bn=batch_norm,
                                  maxpool=True, pool_kernel=2, pool_stride=2, pool_pad=0)

        if self.depthconv:
            self.conv5_1_depthconvweight = 1.  # nn.Parameter(torch.ones(1))
            self.downsample_depth5_1 = nn.AvgPool2d(2, padding=0, stride=2)
            self.conv5_1 = DepthConvModule(512, 512, bn=batch_norm, dilation=2, padding=2)
        else:
            self.conv5_1 = ConvModule(512, 512, bn=batch_norm, dilation=2, padding=2)
        self.conv5_2 = ConvModule(512, 512, bn=batch_norm, dilation=2, padding=2)
        self.conv5_3 = ConvModule(512, 512, bn=batch_norm, dilation=2, padding=2,
                                  maxpool=True, pool_kernel=3, pool_stride=1, pool_pad=1)
        self.pool5a = nn.AvgPool2d(kernel_size=3, stride=1, padding=1)
        self.pool5a_d = DepthAvgPooling(kernel_size=3, stride=1, padding=1)

        self.conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6)  # atrous convolution

        self.conv7 = nn.Conv2d(1024, 1024, kernel_size=1)

    def forward(self, x, depth=None):
        if self.depthconv:
            x = self.conv1_1(x, self.conv1_1_depthconvweight * depth)
        else:
            x = self.conv1_1(x)
        x = self.conv1_2(x)
        if self.depthconv:
            depth = self.downsample_depth2_1(depth)
            x = self.conv2_1(x, self.conv2_1_depthconvweight * depth)
        else:
            x = self.conv2_1(x)
        # print 'xxxxxx',x.size()
        x = self.conv2_2(x)
        if self.depthconv:
            depth = self.downsample_depth3_1(depth)
            x = self.conv3_1(x, self.conv3_1_depthconvweight * depth)
        else:
            x = self.conv3_1(x)
        x = self.conv3_2(x)
        x = self.conv3_3(x)
        if self.depthconv:
            depth = self.downsample_depth4_1(depth)
            # print (depth.mean(),depth.max(),depth.min())
            # torchvision.utils.save_image(depth.data, 'depth.png')
            x = self.conv4_1(x, self.conv4_1_depthconvweight * depth)
        else:
            x = self.conv4_1(x)
        x = self.conv4_2(x)
        conv4_3_feats = x
        x = self.conv4_3(x)
        if self.depthconv:
            depth = self.downsample_depth5_1(depth)
            x = self.conv5_1(x, self.conv5_1_depthconvweight * depth)
        else:
            x = self.conv5_1(x)
        x = self.conv5_2(x)
        x = self.conv5_3(x)
        # x = self.pool5a(x,depth)
        if self.depthconv:
            x = self.pool5a_d(x, depth)
        else:
            x = self.pool5a(x)

        x = F.relu(self.conv6(x))

        conv7_feats = F.relu(self.conv7(x))

        return depth, conv4_3_feats, conv7_feats  # 针对这两个特征层来进行后续的目标检测

def make_layers(cfg, depth_cfg=[], batch_norm=False, depthconv=False):
    layers = []
    in_channels = 3
    for iv, v in enumerate(cfg):
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            if depthconv and iv in depth_cfg:
                conv2d = DepthConv(in_channels, v, kernel_size=3, padding=1)
            else:
                conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)

class AuxiliaryConvolutions(nn.Module):#多尺度预测的网络部分
    def __init__(self, batch_norm=False, depthconv=False):
        super(AuxiliaryConvolutions, self).__init__()

        self.depthconv = depthconv
        if self.depthconv:
            self.conv8_1_depthconvweight = 1.  # nn.Parameter(torch.ones(1))
            # 卷积层输出尺寸： o = ⌊(i + 2p - k) / s⌋ + 1,
            self.conv8_1 = DepthConvModule(1024, 256, kernel_size=1, padding=0, bn=batch_norm)
        else:
            self.conv8_1 = ConvModule(1024, 256, kernel_size=1, padding=0, bn=batch_norm)
        self.conv8_2 = ConvModule(256, 512,kernel_size=3, stride=2, padding=1, bn=batch_norm)

        if self.depthconv:
            self.conv9_1_depthconvweight = 1.  # nn.Parameter(torch.ones(1))
            self.downsample_depth9_1 = nn.AvgPool2d(kernel_size=3, stride=2, padding=1)
            self.conv9_1 = DepthConvModule(512, 128, kernel_size=1, padding=0, bn=batch_norm)
        else:
            self.conv9_1 = ConvModule(512, 128, kernel_size=1, padding=0, bn=batch_norm)
        self.conv9_2 = ConvModule(128, 256,kernel_size=3, stride=2, padding=1, bn=batch_norm)

        if self.depthconv:
            self.conv10_1_depthconvweight = 1.  # nn.Parameter(torch.ones(1))
            self.downsample_depth10_1 = nn.AvgPool2d(kernel_size=3, stride=2, padding=1)
            self.conv10_1 = DepthConvModule(256, 128, kernel_size=1, padding=0, bn=batch_norm)
        else:
            self.conv10_1 = ConvModule(256, 128, kernel_size=1, padding=0, bn=batch_norm)
        self.conv10_2 = ConvModule(128, 256, kernel_size=3, padding=0, bn=batch_norm)

        if self.depthconv:
            self.conv11_1_depthconvweight = 1.  # nn.Parameter(torch.ones(1))
            self.downsample_depth11_1 = nn.AvgPool2d(kernel_size=3, stride=1, padding=0)
            self.conv11_1 = DepthConvModule(256, 128, kernel_size=1, padding=0, bn=batch_norm)
        else:
            self.conv11_1 = ConvModule(256, 128, kernel_size=1, padding=0, bn=batch_norm)
        self.conv11_2 = ConvModule(128, 256, kernel_size=3, padding=0, bn=batch_norm)

    def forward(self, conv7_feats, depth=None):
        if self.depthconv:
            out = self.conv8_1(conv7_feats, self.conv8_1_depthconvweight * depth)  # (N, 256, 19, 19)
        else:
            out = self.conv8_1(conv7_feats)
        out = self.conv8_2(out)  # (N, 512, 10, 10)
        conv8_2_feats = out  # (N, 512, 10, 10)

        if self.depthconv:
            depth = self.downsample_depth9_1(depth)
            out = self.conv9_1(out, self.conv9_1_depthconvweight * depth)
        else:
            out = self.conv9_1(out)
        out = self.conv9_2(out)
        conv9_2_feats = out

        if self.depthconv:
            depth = self.downsample_depth10_1(depth)
            out = self.conv10_1(out, self.conv10_1_depthconvweight * depth)
        else:
            out = self.conv10_1(out)
        out = self.conv10_2(out)
        conv10_2_feats = out

        if self.depthconv:
            depth = self.downsample_depth11_1(depth)
            out = self.conv11_1(out, self.conv11_1_depthconvweight * depth)
        else:
            out = self.conv11_1(out)
        out = self.conv11_2(out)
        conv11_2_feats = out

        return conv8_2_feats, conv9_2_feats, conv10_2_feats, conv11_2_feats  # 这是几个预测层

class PredictionConvolutions(nn.Module):  #根据之前的得到的特征层进行预测
    def __init__(self, n_classes):
        super(PredictionConvolutions, self).__init__()
        self.n_classes = n_classes
        n_boxes = {'conv4_3': 4,  # 之前在特征图上预先产生的先验框的数量与卷积核对应
                   'conv7': 6,
                   'conv8_2': 6,
                   'conv9_2': 6,
                   'conv10_2': 4,
                   'conv11_2': 4}

        #每个坐标有四个参数
        self.loc_conv4_3 = nn.Conv2d(512, n_boxes['conv4_3'] * 4, kernel_size=3, padding=1)
        self.loc_conv7 = nn.Conv2d(1024, n_boxes['conv7'] * 4, kernel_size=3, padding=1)
        self.loc_conv8_2 = nn.Conv2d(512, n_boxes['conv8_2'] * 4, kernel_size=3, padding=1)
        self.loc_conv9_2 = nn.Conv2d(256, n_boxes['conv9_2'] * 4, kernel_size=3, padding=1)
        self.loc_conv10_2 = nn.Conv2d(256, n_boxes['conv10_2'] * 4, kernel_size=3, padding=1)
        self.loc_conv11_2 = nn.Conv2d(256, n_boxes['conv11_2'] * 4, kernel_size=3, padding=1)

        # 计算其中的分类卷积
        self.cl_conv4_3 = nn.Conv2d(512, n_boxes['conv4_3'] * n_classes, kernel_size=3, padding=1)
        self.cl_conv7 = nn.Conv2d(1024, n_boxes['conv7'] * n_classes, kernel_size=3, padding=1)
        self.cl_conv8_2 = nn.Conv2d(512, n_boxes['conv8_2'] * n_classes, kernel_size=3, padding=1)
        self.cl_conv9_2 = nn.Conv2d(256, n_boxes['conv9_2'] * n_classes, kernel_size=3, padding=1)
        self.cl_conv10_2 = nn.Conv2d(256, n_boxes['conv10_2'] * n_classes, kernel_size=3, padding=1)
        self.cl_conv11_2 = nn.Conv2d(256, n_boxes['conv11_2'] * n_classes, kernel_size=3, padding=1)

    def forward(self, conv4_3_feats, conv7_feats, conv8_2_feats, conv9_2_feats, conv10_2_feats, conv11_2_feats):
        batch_size = conv4_3_feats.size(0)

        # 计算其中的框的回归参数
        l_conv4_3 = self.loc_conv4_3(conv4_3_feats)  # (N, 16, 38, 38)
        l_conv4_3 = l_conv4_3.permute(0, 2, 3, 1).contiguous()  # (N, 38, 38, 16), 把通道数放到后面去便于之后的计算
        l_conv4_3 = l_conv4_3.view(batch_size, -1, 4)  # (N, 5776, 4)把每个框放到中间

        l_conv7 = self.loc_conv7(conv7_feats)  # (N, 24, 19, 19)
        l_conv7 = l_conv7.permute(0, 2, 3, 1).contiguous()  # (N, 19, 19, 24)
        l_conv7 = l_conv7.view(batch_size, -1, 4)  # (N, 2166, 4),

        l_conv8_2 = self.loc_conv8_2(conv8_2_feats)  # (N, 24, 10, 10)
        l_conv8_2 = l_conv8_2.permute(0, 2, 3, 1).contiguous()  # (N, 10, 10, 24)
        l_conv8_2 = l_conv8_2.view(batch_size, -1, 4)  # (N, 600, 4)

        l_conv9_2 = self.loc_conv9_2(conv9_2_feats)  # (N, 24, 5, 5)
        l_conv9_2 = l_conv9_2.permute(0, 2, 3, 1).contiguous()  # (N, 5, 5, 24)
        l_conv9_2 = l_conv9_2.view(batch_size, -1, 4)  # (N, 150, 4)

        l_conv10_2 = self.loc_conv10_2(conv10_2_feats)  # (N, 16, 3, 3)
        l_conv10_2 = l_conv10_2.permute(0, 2, 3, 1).contiguous()  # (N, 3, 3, 16)
        l_conv10_2 = l_conv10_2.view(batch_size, -1, 4)  # (N, 36, 4)

        l_conv11_2 = self.loc_conv11_2(conv11_2_feats)  # (N, 16, 1, 1)
        l_conv11_2 = l_conv11_2.permute(0, 2, 3, 1).contiguous()  # (N, 1, 1, 16)
        l_conv11_2 = l_conv11_2.view(batch_size, -1, 4)  # (N, 4, 4)

        # 计算检测的分类
        c_conv4_3 = self.cl_conv4_3(conv4_3_feats)  # (N, 4 * n_classes, 38, 38)
        c_conv4_3 = c_conv4_3.permute(0, 2, 3, 1).contiguous()  # (N, 38, 38, 4 * n_classes)把通道数放到后面去了
        c_conv4_3 = c_conv4_3.view(batch_size, -1, self.n_classes)  # (N, 5776, n_classes)计算每个类别的概率

        c_conv7 = self.cl_conv7(conv7_feats)  # (N, 6 * n_classes, 19, 19)
        c_conv7 = c_conv7.permute(0, 2, 3, 1).contiguous()  # (N, 19, 19, 6 * n_classes)
        c_conv7 = c_conv7.view(batch_size, -1, self.n_classes)  # (N, 2166, n_classes),

        c_conv8_2 = self.cl_conv8_2(conv8_2_feats)  # (N, 6 * n_classes, 10, 10)
        c_conv8_2 = c_conv8_2.permute(0, 2, 3, 1).contiguous()  # (N, 10, 10, 6 * n_classes)
        c_conv8_2 = c_conv8_2.view(batch_size, -1, self.n_classes)  # (N, 600, n_classes)

        c_conv9_2 = self.cl_conv9_2(conv9_2_feats)  # (N, 6 * n_classes, 5, 5)
        c_conv9_2 = c_conv9_2.permute(0, 2, 3, 1).contiguous()  # (N, 5, 5, 6 * n_classes)
        c_conv9_2 = c_conv9_2.view(batch_size, -1, self.n_classes)  # (N, 150, n_classes)

        c_conv10_2 = self.cl_conv10_2(conv10_2_feats)  # (N, 4 * n_classes, 3, 3)
        c_conv10_2 = c_conv10_2.permute(0, 2, 3, 1).contiguous()  # (N, 3, 3, 4 * n_classes)
        c_conv10_2 = c_conv10_2.view(batch_size, -1, self.n_classes)  # (N, 36, n_classes)

        c_conv11_2 = self.cl_conv11_2(conv11_2_feats)  # (N, 4 * n_classes, 1, 1)
        c_conv11_2 = c_conv11_2.permute(0, 2, 3, 1).contiguous()  # (N, 1, 1, 4 * n_classes)
        c_conv11_2 = c_conv11_2.view(batch_size, -1, self.n_classes)  # (N, 4, n_classes)

        # 把所有的边框回归于分类集合在一起
        locs = torch.cat([l_conv4_3, l_conv7, l_conv8_2, l_conv9_2, l_conv10_2, l_conv11_2], dim=1)  # (N, 8732, 4)
        classes_scores = torch.cat([c_conv4_3, c_conv7, c_conv8_2, c_conv9_2, c_conv10_2, c_conv11_2],
                                   dim=1)  # (N, 8732, n_classes)

        return locs, classes_scores  # 返回得到的预测的框于分类的情况

class SSD(nn.Module):
    def __init__(self, n_classes, init_weights=True, depthconv=False,bn=False):
        super(SSD, self).__init__()

        self.n_classes = n_classes

        self.base = VGG_layer(batch_norm=bn,depthconv=depthconv)  # 基础特征提取层,前面两个尺度的网络
        self.aux_convs = AuxiliaryConvolutions(batch_norm=bn,depthconv=depthconv)  # 多尺度辅助预测,后面几个尺度的网络
        self.pred_convs = PredictionConvolutions(n_classes)  # 预测物体的位置以及分类

        if init_weights:
            self._initialize_weights()

        # L2正则化
        # 512通道数在conv4_3_feats,nn.Parameter将torch转为可训练
        self.rescale_factors = nn.Parameter(torch.FloatTensor(1, 512, 1, 1))
        nn.init.constant_(self.rescale_factors, 20)  # 用20填充的张量

        # 在特征图上画框框
        self.priors_cxcy = self.create_prior_boxes()

    def forward(self, image, depth=None):
        # 通过前面的vgg初步提取特征 返回两个尺度的特征层
        depth, conv4_3_feats, conv7_feats = self.base(image, depth)  # (N, 512, 38, 38), (N, 1024, 19, 19)

        # L2正则化
        norm = conv4_3_feats.pow(2).sum(dim=1, keepdim=True).sqrt()  # (N, 1, 38, 38)
        conv4_3_feats = conv4_3_feats / norm  # (N, 512, 38, 38)
        conv4_3_feats = conv4_3_feats * self.rescale_factors  # (N, 512, 38, 38)

        # 后续的特征提取网络
        conv8_2_feats, conv9_2_feats, conv10_2_feats, conv11_2_feats = \
            self.aux_convs(conv7_feats, depth)  # (N, 512, 10, 10),  (N, 256, 5, 5), (N, 256, 3, 3), (N, 256, 1, 1)

        # 利用网路预测几张特征图上的 图像类别以及位置
        locs, classes_scores = self.pred_convs(conv4_3_feats, conv7_feats, conv8_2_feats, conv9_2_feats, conv10_2_feats,
                                               conv11_2_feats)  # (N, 8732, 4), (N, 8732, n_classes)

        return locs, classes_scores

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, sqrt(2. / n))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

    def create_prior_boxes(self):
        # 在特征图上创建先眼眶
        fmap_dims = {'conv4_3': 38,
                     'conv7': 19,
                     'conv8_2': 10,
                     'conv9_2': 5,
                     'conv10_2': 3,
                     'conv11_2': 1}

        obj_scales = {'conv4_3': 0.1,
                      'conv7': 0.2,
                      'conv8_2': 0.375,
                      'conv9_2': 0.55,
                      'conv10_2': 0.725,
                      'conv11_2': 0.9}

        aspect_ratios = {'conv4_3': [1., 2., 0.5],
                         'conv7': [1., 2., 3., 0.5, .333],
                         'conv8_2': [1., 2., 3., 0.5, .333],
                         'conv9_2': [1., 2., 3., 0.5, .333],
                         'conv10_2': [1., 2., 0.5],
                         'conv11_2': [1., 2., 0.5]}

        fmaps = list(fmap_dims.keys())

        prior_boxes = []

        for k, fmap in enumerate(fmaps):
            for i in range(fmap_dims[fmap]):
                for j in range(fmap_dims[fmap]):
                    cx = (j + 0.5) / fmap_dims[fmap]
                    cy = (i + 0.5) / fmap_dims[fmap]

                    for ratio in aspect_ratios[fmap]:
                        prior_boxes.append([cx, cy, obj_scales[fmap] * sqrt(ratio), obj_scales[fmap] / sqrt(ratio)])

                        if ratio == 1.:
                            try:
                                additional_scale = sqrt(obj_scales[fmap] * obj_scales[fmaps[k + 1]])
                            except IndexError:
                                additional_scale = 1.
                            prior_boxes.append([cx, cy, additional_scale, additional_scale])

        prior_boxes = torch.FloatTensor(prior_boxes).to(device)  # (8732, 4)
        prior_boxes.clamp_(0, 1)  # (8732, 4) 裁剪多出边界的框

        return prior_boxes

    def detect_objects(self, predicted_locs, predicted_scores, min_score, max_overlap, top_k):
        # 检测物体
        batch_size = predicted_locs.size(0)  # (N, 8732, 4)
        n_priors = self.priors_cxcy.size(0)   # (8732, 4) 检测框
        predicted_scores = F.softmax(predicted_scores, dim=2)  # (N, 8732, n_classes) 计算分类的分数

        # 将所有信息的载体先创建表格来存放
        all_images_boxes = list()
        all_images_labels = list()
        all_images_scores = list()

        assert n_priors == predicted_locs.size(1) == predicted_scores.size(1)  # 判断数量正确与否

        for i in range(batch_size):
            # 从之前的计算的数值中解码
            decoded_locs = cxcy_to_xy(
                gcxgcy_to_cxcy(predicted_locs[i], self.priors_cxcy))  # (8732, 4), 计算出坐标

            # 创建列表存储信息
            image_boxes = list()
            image_labels = list()
            image_scores = list()

            max_scores, best_label = predicted_scores[i].max(dim=1)  # (8732)

            for c in range(1, self.n_classes):

                class_scores = predicted_scores[i][:, c]  # (8732)
                score_above_min_score = class_scores > min_score  # 大于分数的门限
                n_above_min_score = score_above_min_score.sum().item()
                if n_above_min_score == 0:
                    continue
                class_scores = class_scores[score_above_min_score]
                class_decoded_locs = decoded_locs[score_above_min_score]

                class_scores, sort_ind = class_scores.sort(dim=0, descending=True)
                class_decoded_locs = class_decoded_locs[sort_ind]

                # (n_qualified, n_min_score) (8732, 8732)中剩下的大于分数的
                overlap = find_jaccard_overlap(class_decoded_locs, class_decoded_locs)

                # (NMS)非极大抑制的过程
                suppress = torch.zeros((n_above_min_score), dtype=torch.uint8).to(device)  # (n_qualified)(8732)中剩下的

                for box in range(class_decoded_locs.size(0)):  # (8732, 4)中剩下的

                    if suppress[box] == 1:
                        continue
                    suppress = torch.max(suppress, overlap[box] > max_overlap)  # 输出最大值
                    suppress[box] = 0

                image_boxes.append(class_decoded_locs[1 - suppress])
                image_labels.append(torch.LongTensor((1 - suppress).sum().item() * [c]).to(device))
                image_scores.append(class_scores[1 - suppress])

            if len(image_boxes) == 0:
                image_boxes.append(torch.FloatTensor([[0., 0., 1., 1.]]).to(device))
                image_labels.append(torch.LongTensor([0]).to(device))
                image_scores.append(torch.FloatTensor([0.]).to(device))

            image_boxes = torch.cat(image_boxes, dim=0)  # (n_objects, 4)
            image_labels = torch.cat(image_labels, dim=0)  # (n_objects)
            image_scores = torch.cat(image_scores, dim=0)  # (n_objects)
            n_objects = image_scores.size(0)

            if n_objects > top_k:
                image_scores, sort_ind = image_scores.sort(dim=0, descending=True)
                image_scores = image_scores[:top_k]  # (top_k)
                image_boxes = image_boxes[sort_ind][:top_k]  # (top_k, 4)
                image_labels = image_labels[sort_ind][:top_k]  # (top_k)

            all_images_boxes.append(image_boxes)
            all_images_labels.append(image_labels)
            all_images_scores.append(image_scores)

        return all_images_boxes, all_images_labels, all_images_scores

class MultiBoxLoss(nn.Module):

    def __init__(self, priors_cxcy, threshold=0.5, neg_pos_ratio=3, alpha=1.):  # priors_cxcy先验框
        super(MultiBoxLoss, self).__init__()
        self.priors_cxcy = priors_cxcy
        self.priors_xy = cxcy_to_xy(priors_cxcy)
        self.threshold = threshold  # 判断门限
        self.neg_pos_ratio = neg_pos_ratio
        self.alpha = alpha

        self.smooth_l1 = nn.L1Loss()  # 光滑标签来计算框的回归损失
        self.cross_entropy = nn.CrossEntropyLoss(reduce=False)  # 交叉熵计算分类损失

    def forward(self, predicted_locs, predicted_scores, boxes, labels):  # 预测框、预测分数、实际框、实际种类
        batch_size = predicted_locs.size(0)  # (N, x, 4) x是8732中剩下的
        n_priors = self.priors_cxcy.size(0)  # (x, 4)
        n_classes = predicted_scores.size(2) # (N, x, nc)

        assert n_priors == predicted_locs.size(1) == predicted_scores.size(1)

        true_locs = torch.zeros((batch_size, n_priors, 4), dtype=torch.float).to(device)  # (N, 8732, 4)
        true_classes = torch.zeros((batch_size, n_priors), dtype=torch.long).to(device)  # (N, 8732)

        for i in range(batch_size):
            n_objects = boxes[i].size(0)

            overlap = find_jaccard_overlap(boxes[i],
                                           self.priors_xy)  # (n_objects, 8732)

            # 根据真实框找到最大的iou的框 torch.max返回最大值和最大值索引
            overlap_for_each_prior, object_for_each_prior = overlap.max(dim=0)  # (8732)

            _, prior_for_each_object = overlap.max(dim=1)  # (N_o)

            object_for_each_prior[prior_for_each_object] = torch.LongTensor(range(n_objects)).to(device)

            overlap_for_each_prior[prior_for_each_object] = 1.

            label_for_each_prior = labels[i][object_for_each_prior]  # (8732)

            label_for_each_prior[overlap_for_each_prior < self.threshold] = 0  # (8732)

            true_classes[i] = label_for_each_prior

            # 对预测框进行编码
            true_locs[i] = cxcy_to_gcxgcy(xy_to_cxcy(boxes[i][object_for_each_prior]), self.priors_cxcy)  # (8732, 4)

        # 正样本的判断
        positive_priors = true_classes != 0  # (N, 8732)

        # 计算位置的回归损失
        loc_loss = self.smooth_l1(predicted_locs[positive_priors], true_locs[positive_priors])  # (), scalar

        n_positives = positive_priors.sum(dim=1)  # (N)
        n_hard_negatives = self.neg_pos_ratio * n_positives  # (N)

        # 计算全部的损失
        conf_loss_all = self.cross_entropy(predicted_scores.view(-1, n_classes), true_classes.view(-1))  # (N * 8732)
        conf_loss_all = conf_loss_all.view(batch_size, n_priors)  # (N, 8732)

        conf_loss_pos = conf_loss_all[positive_priors]  # (sum(n_positives))

        # 找难的负样本
        # 存储负样本
        conf_loss_neg = conf_loss_all.clone()  # (N, 8732)
        conf_loss_neg[positive_priors] = 0.  # (N, 8732),
        conf_loss_neg, _ = conf_loss_neg.sort(dim=1, descending=True)  # (N, 8732)
        hardness_ranks = torch.LongTensor(range(n_priors)).unsqueeze(0).expand_as(conf_loss_neg).to(device)  # (N, 8732)
        hard_negatives = hardness_ranks < n_hard_negatives.unsqueeze(1)  # (N, 8732)
        conf_loss_hard_neg = conf_loss_neg[hard_negatives]  # (sum(n_hard_negatives))

        conf_loss = (conf_loss_hard_neg.sum() + conf_loss_pos.sum()) / n_positives.sum().float()  # (), scalar

        return conf_loss + self.alpha * loc_loss  # 最后的总损失